import torch.nn
from torch import Tensor
from torchvision.models.resnet import BasicBlock, ResNet


class ResNet18(ResNet):
    def __init__(self):
        super(ResNet18, self).__init__(BasicBlock, [2, 2, 2, 2])
        self.client_feature_extraction = torch.nn.Sequential(
            self.conv1,
            self.bn1,
            self.relu,
            self.maxpool,
        )
        self.server_feature_extraction = torch.nn.Sequential(
            self.layer1,
            self.layer2,
            self.layer3,
            self.layer4,
            self.avgpool,
        )
        self.server_classifier = self.fc

    def __str__(self):
        return "ResNet18"

    def forward(self, x: Tensor) -> Tensor:
        x = self.client_feature_extraction(x)
        x = self.server_feature_extraction(x)
        x = self.server_classifier(x.view(x.size(0), -1))
        return x


class ClientSideResNet18(ResNet18):
    def __init__(self):
        super(ClientSideResNet18, self).__init__()
        del self.server_feature_extraction
        del self.server_classifier

    def forward(self, x):
        return self.client_feature_extraction(x)


class ServerSideResNet18(ResNet18):
    def __init__(self):
        super(ServerSideResNet18, self).__init__()
        del self.client_feature_extraction

    def forward(self, x):
        x = self.server_feature_extraction(x)
        return self.server_classifier(x.view(x.size(0), -1))
